# translate.py
# 调用方法
# from translate import translate_text
# translated_text = translate_text("你好，世界！")

from transformers import pipeline
from langdetect import detect

# Initialize translation pipelines
model_name_zh_to_en = "./machine_translation/trans_model_zh-en"
model_name_en_to_zh = "./machine_translation/trans_model_en-zh"

translation_pipeline_zh_to_en = pipeline(
    "translation", model=model_name_zh_to_en, tokenizer=model_name_zh_to_en)
translation_pipeline_en_to_zh = pipeline(
    "translation", model=model_name_en_to_zh, tokenizer=model_name_en_to_zh)


def detect_language(input_text):
    """Detect the language of the input text."""
    return detect(input_text)


def translate_text(input_text):
    """Automatically detect language and translate text."""
    language = detect_language(input_text)

    if language == 'zh-cn' or language == 'zh':
        # Translate Chinese to English
        translate_result = translation_pipeline_zh_to_en(
            input_text, max_length=400)
    elif language == 'en':
        # Translate English to Chinese
        translate_result = translation_pipeline_en_to_zh(
            input_text, max_length=4000)
    else:
        return "不支持中文EN混译~"

    translated_text = translate_result[0]['translation_text']
    return translated_text
